import torch
from torch import nn
import numpy as np
class NeuralNetwork(nn.Module):
  def __init__(self, layers, device):

    super().__init__()
    self.layers = layers
    self.linears = nn.ModuleList([nn.Linear(self.layers[i], self.layers[i+1], device=device) for i in range(len(self.layers)-1)])
    self.activation=nn.Tanh()

    for i in range(len(self.layers)-1):
      nn.init.xavier_normal_(self.linears[i].weight.data, gain=1.0)
      nn.init.zeros_(self.linears[i].bias.data)

  def forward(self, x, t=None):
    if t is not None:
      input = torch.cat((x, t), dim=1)
    else:
      input = x
    for i in range(len(self.layers)-2):
        z = self.linears[i](input)
        input = self.activation(z)

    input = self.linears[-1](input)
    return input